import numpy as np
import scipy.io.wavfile
import scipy.signal
import matplotlib.pyplot as plt
import seaborn as sns
import os

VER = 'v1.0'
BASE_DIR, FILE_NAME = os.path.split(__file__)
path = '../../../../../large_data/audio/zsn-stop1.wav'
AUDIO_PATH = os.path.join(BASE_DIR, path)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)

sr, signal = scipy.io.wavfile.read(AUDIO_PATH)
if len(signal.shape) >= 2:
    signal = signal[:, 0]

nperseg = int(sr / 1000. * 25)
noverlap = nperseg - int(sr / 1000. * 10)
print('nperseg', nperseg)
print('noverlap', noverlap)
_, _, spectrum = scipy.signal.spectrogram(
    signal,
    fs=sr,
    window='hann',
    nperseg=nperseg,
    noverlap=noverlap,
    detrend=None
)
spectrum = np.float32(spectrum)
spectrum = np.where(spectrum <= 0., np.finfo(np.float32).eps, spectrum)
spectrum = np.log10(spectrum)
sns.heatmap(spectrum)
plt.show()
